#include <mcp/auth/auth_handler.h>
#include <mcp/error.h>
#include <mcp/json/value.h>

#include <iostream>
#include <sstream>
#include <chrono>
#include <memory>
#include <algorithm>
#include <regex>

namespace mcp {
namespace auth {

// AuthHandler implementation
struct AuthHandler::Impl {
    std::unordered_map<std::string, std::unique_ptr<Authenticator>> authenticators;
    std::string default_scheme;
    bool authentication_required = false;
    std::function<void(const Error&)> failure_handler;
    mutable std::mutex handler_mutex;

    Impl() = default;
};

AuthHandler::AuthHandler() : impl_(std::make_unique<Impl>()) {}

AuthHandler::~AuthHandler() = default;

void AuthHandler::add_authenticator(std::unique_ptr<Authenticator> authenticator) {
    if (!authenticator) {
        return;
    }

    std::lock_guard<std::mutex> lock(impl_->handler_mutex);
    std::string scheme = authenticator->get_scheme();
    impl_->authenticators[scheme] = std::move(authenticator);

    // Set as default if it's the first authenticator
    if (impl_->default_scheme.empty()) {
        impl_->default_scheme = scheme;
    }
}

void AuthHandler::remove_authenticator(const std::string& scheme) {
    std::lock_guard<std::mutex> lock(impl_->handler_mutex);
    impl_->authenticators.erase(scheme);

    // Update default scheme if needed
    if (impl_->default_scheme == scheme) {
        impl_->default_scheme = impl_->authenticators.empty() ? "" : impl_->authenticators.begin()->first;
    }
}

Result<AuthResult> AuthHandler::authenticate(const std::string& auth_header) {
    if (auth_header.empty()) {
        return Error::validation_error("Empty authorization header");
    }

    // Parse authorization header
    std::string scheme;
    std::string credentials;

    // Extract scheme and credentials
    size_t space_pos = auth_header.find(' ');
    if (space_pos == std::string::npos) {
        return Error::validation_error("Invalid authorization header format");
    }

    scheme = auth_header.substr(0, space_pos);
    credentials = auth_header.substr(space_pos + 1);

    // Find authenticator for scheme
    std::lock_guard<std::mutex> lock(impl_->handler_mutex);
    auto it = impl_->authenticators.find(scheme);
    if (it == impl_->authenticators.end()) {
        return Error::validation_error("Unsupported authentication scheme: " + scheme);
    }

    // Create credentials object
    AuthCredentials auth_creds;
    if (scheme == "Basic") {
        // Parse Basic auth credentials
        // In a real implementation, this would decode base64
        auth_creds.username = credentials; // Simplified
        auth_creds.password = ""; // Would be extracted from decoded credentials
    } else if (scheme == "Bearer") {
        auth_creds.bearer_token = credentials;
    } else if (scheme == "ApiKey") {
        // Try to determine header name from authenticator
        auto* api_key_auth = dynamic_cast<ApiKeyAuthenticator*>(it->second.get());
        if (api_key_auth) {
            auth_creds.api_key = credentials;
        }
    }

    // Authenticate using the appropriate authenticator
    auto result = it->second->authenticate(auth_creds);
    if (!result) {
        if (impl_->failure_handler) {
            impl_->failure_handler(result.error());
        }
        return result;
    }

    return result;
}

Result<AuthResult> AuthHandler::authenticate(const Credentials& credentials) {
    if (impl_->authenticators.empty()) {
        return Error::validation_error("No authenticators configured");
    }

    // Try default scheme first
    if (!impl_->default_scheme.empty()) {
        auto it = impl_->authenticators.find(impl_->default_scheme);
        if (it != impl_->authenticators.end()) {
            auto result = it->second->authenticate(credentials);
            if (result) {
                return result;
            }
        }
    }

    // Try all authenticators
    for (auto& [scheme, authenticator] : impl_->authenticators) {
        auto result = authenticator->authenticate(credentials);
        if (result) {
            return result;
        }
    }

    return Error::authentication_error("Authentication failed with all configured schemes");
}

std::string AuthHandler::generate_challenge() const {
    std::lock_guard<std::mutex> lock(impl_->handler_mutex);

    if (impl_->authenticators.empty()) {
        return "";
    }

    std::vector<AuthChallenge> challenges;
    for (const auto& [scheme, authenticator] : impl_->authenticators) {
        auto scheme_challenges = authenticator->get_challenges();
        challenges.insert(challenges.end(), scheme_challenges.begin(), scheme_challenges.end());
    }

    return generate_www_authenticate_header(challenges);
}

bool AuthHandler::validate_token(const std::string& token) const {
    if (token.empty()) {
        return false;
    }

    std::lock_guard<std::mutex> lock(impl_->handler_mutex);

    // Try to determine token type and validate accordingly
    if (token.substr(0, 6) == "Basic ") {
        auto it = impl_->authenticators.find("Basic");
        if (it != impl_->authenticators.end()) {
            return it->second->validate_token(token);
        }
    } else if (token.substr(0, 7) == "Bearer ") {
        auto it = impl_->authenticators.find("Bearer");
        if (it != impl_->authenticators.end()) {
            return it->second->validate_token(token);
        }
    } else {
        // Try all authenticators for generic token validation
        for (const auto& [scheme, authenticator] : impl_->authenticators) {
            if (authenticator->validate_token(token)) {
                return true;
            }
        }
    }

    return false;
}

void AuthHandler::set_default_scheme(const std::string& scheme) {
    std::lock_guard<std::mutex> lock(impl_->handler_mutex);
    if (impl_->authenticators.find(scheme) != impl_->authenticators.end()) {
        impl_->default_scheme = scheme;
    }
}

const std::string& AuthHandler::get_default_scheme() const {
    std::lock_guard<std::mutex> lock(impl_->handler_mutex);
    return impl_->default_scheme;
}

std::vector<std::string> AuthHandler::get_supported_schemes() const {
    std::lock_guard<std::mutex> lock(impl_->handler_mutex);

    std::vector<std::string> schemes;
    for (const auto& [scheme, _] : impl_->authenticators) {
        schemes.push_back(scheme);
    }
    return schemes;
}

bool AuthHandler::is_authentication_required() const {
    return impl_->authentication_required;
}

void AuthHandler::set_authentication_required(bool required) {
    impl_->authentication_required = required;
}

void AuthHandler::set_failure_handler(std::function<void(const Error&)> handler) {
    impl_->failure_handler = std::move(handler);
}

// TokenManager implementation
struct TokenManager::Impl {
    std::string access_token;
    std::string refresh_token;
    std::chrono::system_clock::time_point expires_at;
    bool token_valid = false;
    std::function<std::string(const std::string&)> refresh_callback;
    mutable std::mutex token_mutex;

    Impl() = default;
};

TokenManager::TokenManager() : impl_(std::make_unique<Impl>()) {}

TokenManager::~TokenManager() = default;

void TokenManager::store_access_token(const std::string& token, int expires_in) {
    std::lock_guard<std::mutex> lock(impl_->token_mutex);
    impl_->access_token = token;
    impl_->token_valid = !token.empty();

    if (expires_in > 0) {
        impl_->expires_at = std::chrono::system_clock::now() + std::chrono::seconds(expires_in);
    } else {
        impl_->expires_at = std::chrono::system_clock::time_point{}; // Never expires
    }
}

void TokenManager::store_refresh_token(const std::string& token) {
    std::lock_guard<std::mutex> lock(impl_->token_mutex);
    impl_->refresh_token = token;
}

std::string TokenManager::get_access_token() const {
    std::lock_guard<std::mutex> lock(impl_->token_mutex);

    // Check if token has expired
    if (impl_->token_valid && impl_->expires_at != std::chrono::system_clock::time_point{}) {
        if (std::chrono::system_clock::now() >= impl_->expires_at) {
            impl_->token_valid = false;
        }
    }

    return impl_->token_valid ? impl_->access_token : "";
}

std::string TokenManager::get_refresh_token() const {
    std::lock_guard<std::mutex> lock(impl_->token_mutex);
    return impl_->refresh_token;
}

bool TokenManager::is_access_token_valid() const {
    std::lock_guard<std::mutex> lock(impl_->token_mutex);

    if (!impl_->token_valid) {
        return false;
    }

    // Check expiration
    if (impl_->expires_at != std::chrono::system_clock::time_point{}) {
        return std::chrono::system_clock::now() < impl_->expires_at;
    }

    return true;
}

bool TokenManager::has_refresh_token() const {
    std::lock_guard<std::mutex> lock(impl_->token_mutex);
    return !impl_->refresh_token.empty();
}

void TokenManager::clear_tokens() {
    std::lock_guard<std::mutex> lock(impl_->token_mutex);
    impl_->access_token.clear();
    impl_->refresh_token.clear();
    impl_->token_valid = false;
    impl_->expires_at = std::chrono::system_clock::time_point{};
}

bool TokenManager::load_tokens(std::function<std::string(const std::string&)> storage) {
    if (!storage) {
        return false;
    }

    try {
        std::lock_guard<std::mutex> lock(impl_->token_mutex);
        impl_->access_token = storage("access_token");
        impl_->refresh_token = storage("refresh_token");

        // Parse expiration time
        std::string expires_str = storage("expires_at");
        if (!expires_str.empty()) {
            std::istringstream iss(expires_str);
            std::chrono::system_clock::time_point::rep timestamp;
            if (iss >> timestamp) {
                impl_->expires_at = std::chrono::system_clock::time_point(std::chrono::system_clock::duration(timestamp));
            }
        }

        impl_->token_valid = !impl_->access_token.empty();
        return true;
    } catch (...) {
        return false;
    }
}

bool TokenManager::save_tokens(std::function<void(const std::string&, const std::string&)> storage) const {
    if (!storage) {
        return false;
    }

    try {
        std::lock_guard<std::mutex> lock(impl_->token_mutex);
        storage("access_token", impl_->access_token);
        storage("refresh_token", impl_->refresh_token);

        // Save expiration time
        std::ostringstream expires_oss;
        if (impl_->expires_at != std::chrono::system_clock::time_point{}) {
            expires_oss << impl_->expires_at.time_since_epoch().count();
        }
        storage("expires_at", expires_oss.str());

        return true;
    } catch (...) {
        return false;
    }
}

void TokenManager::set_refresh_callback(std::function<std::string(const std::string&)> callback) {
    std::lock_guard<std::mutex> lock(impl_->token_mutex);
    impl_->refresh_callback = std::move(callback);
}

bool TokenManager::refresh_access_token() {
    std::lock_guard<std::mutex> lock(impl_->token_mutex);

    if (!impl_->refresh_callback || impl_->refresh_token.empty()) {
        return false;
    }

    try {
        std::string new_token = impl_->refresh_callback(impl_->refresh_token);
        if (!new_token.empty()) {
            impl_->access_token = new_token;
            impl_->token_valid = true;
            // Reset expiration - caller should update this
            return true;
        }
    } catch (...) {
        // Ignore callback errors
    }

    return false;
}

// AuthorizationHandler implementation
struct AuthorizationHandler::Impl {
    std::unordered_map<std::string, std::vector<std::string>> role_permissions;
    std::unordered_map<std::string, std::vector<std::string>> user_permissions;
    mutable std::mutex auth_mutex;

    Impl() = default;
};

AuthorizationHandler::AuthorizationHandler() : impl_(std::make_unique<Impl>()) {}

AuthorizationHandler::~AuthorizationHandler() = default;

bool AuthorizationHandler::has_permission(const AuthContext& context, const std::string& permission) const {
    if (!context.authenticated) {
        return false;
    }

    std::lock_guard<std::mutex> lock(impl_->auth_mutex);

    // Check user-specific permissions first
    auto user_it = impl_->user_permissions.find(context.user_id);
    if (user_it != impl_->user_permissions.end()) {
        const auto& permissions = user_it->second;
        if (std::find(permissions.begin(), permissions.end(), permission) != permissions.end()) {
            return true;
        }
    }

    // Check role-based permissions
    for (const auto& role : context.roles) {
        auto role_it = impl_->role_permissions.find(role);
        if (role_it != impl_->role_permissions.end()) {
            const auto& permissions = role_it->second;
            if (std::find(permissions.begin(), permissions.end(), permission) != permissions.end()) {
                return true;
            }
        }
    }

    return false;
}

bool AuthorizationHandler::has_role(const AuthContext& context, const std::string& role) const {
    if (!context.authenticated) {
        return false;
    }

    return std::find(context.roles.begin(), context.roles.end(), role) != context.roles.end();
}

bool AuthorizationHandler::has_any_role(const AuthContext& context, const std::vector<std::string>& roles) const {
    if (!context.authenticated) {
        return false;
    }

    for (const auto& role : roles) {
        if (std::find(context.roles.begin(), context.roles.end(), role) != context.roles.end()) {
            return true;
        }
    }

    return false;
}

bool AuthorizationHandler::has_all_roles(const AuthContext& context, const std::vector<std::string>& roles) const {
    if (!context.authenticated) {
        return false;
    }

    for (const auto& role : roles) {
        if (std::find(context.roles.begin(), context.roles.end(), role) == context.roles.end()) {
            return false;
        }
    }

    return true;
}

void AuthorizationHandler::add_role_permissions(const std::string& role, const std::vector<std::string>& permissions) {
    std::lock_guard<std::mutex> lock(impl_->auth_mutex);
    impl_->role_permissions[role] = permissions;
}

void AuthorizationHandler::add_user_permissions(const std::string& user_id, const std::vector<std::string>& permissions) {
    std::lock_guard<std::mutex> lock(impl_->auth_mutex);
    impl_->user_permissions[user_id] = permissions;
}

void AuthorizationHandler::clear_rules() {
    std::lock_guard<std::mutex> lock(impl_->auth_mutex);
    impl_->role_permissions.clear();
    impl_->user_permissions.clear();
}

// Utility functions
std::string auth_context_to_string(const AuthContext& context) {
    std::ostringstream ss;
    ss << "AuthContext{";
    ss << "authenticated=" << (context.authenticated ? "true" : "false");

    if (context.authenticated) {
        ss << ", user_id=" << context.user_id;
        ss << ", username=" << context.username;

        if (!context.roles.empty()) {
            ss << ", roles=[";
            bool first = true;
            for (const auto& role : context.roles) {
                if (!first) ss << ", ";
                ss << role;
                first = false;
            }
            ss << "]";
        }

        if (!context.claims.empty()) {
            ss << ", claims=" << context.claims.size() << " items";
        }
    }

    ss << "}";
    return ss.str();
}

bool validate_auth_context(const AuthContext& context) {
    if (!context.authenticated) {
        // Unauthenticated contexts should have minimal valid data
        return context.user_id.empty() &&
               context.username.empty() &&
               context.roles.empty() &&
               context.access_token.empty();
    }

    // Authenticated contexts must have required fields
    if (context.user_id.empty() || context.access_token.empty()) {
        return false;
    }

    // Validate expiration
    if (context.expires_at != std::chrono::system_clock::time_point{}) {
        if (std::chrono::system_clock::now() >= context.expires_at) {
            return false; // Token expired
        }
    }

    return true;
}

AuthContext create_auth_context_from_token(const std::string& access_token,
                                          const std::string& user_id,
                                          const std::string& username,
                                          const std::vector<std::string>& roles,
                                          const std::unordered_map<std::string, std::string>& claims,
                                          int expires_in) {
    AuthContext context;
    context.authenticated = true;
    context.access_token = access_token;
    context.user_id = user_id;
    context.username = username;
    context.roles = roles;
    context.claims = claims;

    if (expires_in > 0) {
        context.expires_at = std::chrono::system_clock::now() + std::chrono::seconds(expires_in);
    }

    return context;
}

} // namespace auth
} // namespace mcp